from numpy import (
    ndarray,  # 数组类型
    finfo,  # 数值极限
    eye,  # 单位矩阵
    matmul,  # 矩阵乘法
    trace,  # 矩阵迹
    sqrt,  # 计算平方根
)
from numpy.linalg import (
    inv,  # 矩阵求逆
)

from src.utils.norm_precoder import (
    norm_precoder,  # 归一化预编码矩阵函数
)


def mmse_precoder_normalized(
        channel_matrix,  # 信道矩阵，表示用户和卫星天线之间的信道状态
        noise_power_watt: float,  # 噪声功率，单位为瓦特
        power_constraint_watt: float,  # 功率约束，单位为瓦特
        sat_nr,  # 卫星数量
        sat_ant_nr,  # 每个卫星的天线数量
) -> ndarray:
    """
    计算归一化的最小均方误差（MMSE）预编码矩阵。

    参数:
    - channel_matrix: 信道矩阵，大小为 (用户数量, 卫星天线数量)
    - noise_power_watt: 噪声功率，单位为瓦特
    - power_constraint_watt: 功率约束，单位为瓦特
    - sat_nr: 卫星数量
    - sat_ant_nr: 每个卫星的天线数量

    返回:
    - ndarray: 归一化的预编码矩阵
    """
    # 计算未归一化的MMSE预编码矩阵
    precoding_matrix = mmse_precoder_no_norm(
        channel_matrix=channel_matrix,
        noise_power_watt=noise_power_watt,
        power_constraint_watt=power_constraint_watt,
    )

    # 对预编码矩阵进行归一化
    precoding_matrix_normed = norm_precoder(
        precoding_matrix=precoding_matrix,
        power_constraint_watt=power_constraint_watt,
        per_satellite=True,
        sat_nr=sat_nr,
        sat_ant_nr=sat_ant_nr,
    )

    return precoding_matrix_normed


def mmse_precoder_no_norm(
        channel_matrix,  # 信道矩阵，表示用户和卫星天线之间的信道状态
        noise_power_watt: float,  # 噪声功率，单位为瓦特
        power_constraint_watt: float,  # 功率约束，单位为瓦特
) -> ndarray:
    """
    计算未归一化的最小均方误差（MMSE）预编码矩阵。

    参数:
    - channel_matrix: 信道矩阵，大小为 (用户数量, 卫星天线数量)
    - noise_power_watt: 噪声功率，单位为瓦特
    - power_constraint_watt: 功率约束，单位为瓦特

    返回:
    - ndarray: 未归一化的预编码矩阵
    """
    # 设定用于矩阵求逆时的常量，避免逆矩阵时出现数值不稳定（这里设为0，实际应用中可以设置为非常小的正数）
    inversion_constant_lambda = 0

    user_nr = channel_matrix.shape[0]  # 用户数量
    sat_tot_ant_nr = channel_matrix.shape[1]  # 总天线数量（所有卫星）

    # 计算MMSE预编码矩阵
    precoding_matrix = (
        matmul(
            inv(
                matmul(channel_matrix.conj().T, channel_matrix)  # 信道矩阵的共轭转置与信道矩阵相乘
                + (noise_power_watt * user_nr / power_constraint_watt + inversion_constant_lambda) * eye(sat_tot_ant_nr)  # 加上噪声功率和常量乘以单位矩阵
            ),
            channel_matrix.conj().T  # 与信道矩阵的共轭转置相乘
        )
    )

    return precoding_matrix
